"""
批量验证API测试脚本

读取verify.csv数据集,调用API进行验证,生成带有预测结果的CSV文件
支持并发请求以提高速度
"""
import pandas as pd
import requests
from tqdm import tqdm
import time
from datetime import datetime
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock


def call_safety_api(text: str, threshold: float = 0.5, api_url: str = "http://localhost:8000/api/validate") -> dict:
    """
    调用安全检测API

    参数:
        text: 待检测文本
        threshold: 置信度阈值
        api_url: API地址

    返回:
        dict: API响应结果
    """
    try:
        response = requests.post(
            api_url,
            json={
                "text": text,
                "threshold": threshold
            },
            timeout=30
        )

        if response.status_code == 200:
            return {
                "status": "success",
                "data": response.json()
            }
        else:
            # 规则层拦截会返回500错误
            error_detail = response.json().get("detail", "未知错误")
            return {
                "status": "rule_blocked",
                "error": error_detail
            }
    except requests.exceptions.Timeout:
        return {
            "status": "timeout",
            "error": "请求超时"
        }
    except Exception as e:
        return {
            "status": "error",
            "error": str(e)
        }


def process_single_item(args):
    """
    处理单条数据(用于并发)

    参数:
        args: (idx, text, true_label, threshold, api_url)

    返回:
        dict: 处理结果
    """
    idx, text, true_label, threshold, api_url = args

    result = call_safety_api(text, threshold, api_url)

    # 构建结果字典
    result_dict = {
        'idx': idx,
        'predicted_label': -1,
        'is_safe': None,
        'confidence': 0.0,
        'api_status': '',
        'error_message': '',
        'correct': False
    }

    if result['status'] == 'success':
        data = result['data']
        is_safe = data.get('is_safe', True)
        predicted_label = 0 if is_safe else 1

        result_dict['predicted_label'] = predicted_label
        result_dict['is_safe'] = is_safe
        result_dict['confidence'] = data.get('confidence', 0.0)
        result_dict['api_status'] = 'success'
        result_dict['correct'] = (predicted_label == true_label)

    elif result['status'] == 'rule_blocked':
        result_dict['predicted_label'] = 1
        result_dict['is_safe'] = False
        result_dict['confidence'] = 1.0
        result_dict['api_status'] = 'rule_blocked'
        result_dict['error_message'] = result.get('error', '')
        result_dict['correct'] = (1 == true_label)

    else:
        result_dict['api_status'] = result['status']
        result_dict['error_message'] = result.get('error', '')
        result_dict['correct'] = False

    return result_dict


def batch_verify_dataset(
    input_csv: str = "data/verify.csv",
    output_csv: str = "results/verify_api_results.csv",
    api_url: str = "http://localhost:8000/api/validate",
    threshold: float = 0.5,
    max_samples: int = None,
    max_workers: int = 10
):
    """
    批量验证数据集(并发版本)

    参数:
        input_csv: 输入CSV文件路径
        output_csv: 输出CSV文件路径
        api_url: API地址
        threshold: 置信度阈值
        max_samples: 最大测试样本数(None表示全部)
        max_workers: 并发线程数
    """
    print("="*70)
    print("📊 批量API验证测试 (并发模式)")
    print("="*70)

    # 1. 读取数据集
    # 转换为绝对路径
    if not os.path.isabs(input_csv):
        script_dir = os.path.dirname(os.path.abspath(__file__))
        input_csv = os.path.join(script_dir, input_csv)

    if not os.path.isabs(output_csv):
        script_dir = os.path.dirname(os.path.abspath(__file__))
        output_csv = os.path.join(script_dir, output_csv)

    print(f"\n📂 读取数据集: {input_csv}")
    df = pd.read_csv(input_csv)

    if max_samples is not None:
        df = df.head(max_samples)
        print(f"⚠️  限制测试样本数: {max_samples}")

    total_samples = len(df)
    print(f"✅ 共 {total_samples} 条数据")
    print(f"⚡ 并发线程数: {max_workers}")

    # 2. 初始化结果列
    df['predicted_label'] = -1
    df['is_safe'] = None
    df['confidence'] = 0.0
    df['api_status'] = ""
    df['error_message'] = ""
    df['correct'] = False

    # 3. 准备任务列表
    tasks = [
        (idx, str(row['text']), int(row['label']), threshold, api_url)
        for idx, row in df.iterrows()
    ]

    # 4. 并发调用API
    print(f"\n🚀 开始并发调用API (阈值={threshold})...")
    print(f"📡 API地址: {api_url}\n")

    start_time = time.time()
    success_count = 0
    rule_blocked_count = 0
    error_count = 0

    # 使用线程池并发处理
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_idx = {executor.submit(process_single_item, task): task[0] for task in tasks}

        # 使用进度条显示进度
        with tqdm(total=total_samples, desc="验证中") as pbar:
            for future in as_completed(future_to_idx):
                result = future.result()
                idx = result['idx']

                # 更新DataFrame
                df.at[idx, 'predicted_label'] = result['predicted_label']
                df.at[idx, 'is_safe'] = result['is_safe']
                df.at[idx, 'confidence'] = result['confidence']
                df.at[idx, 'api_status'] = result['api_status']
                df.at[idx, 'error_message'] = result['error_message']
                df.at[idx, 'correct'] = result['correct']

                # 统计
                if result['api_status'] == 'success':
                    success_count += 1
                elif result['api_status'] == 'rule_blocked':
                    rule_blocked_count += 1
                else:
                    error_count += 1

                pbar.update(1)

    end_time = time.time()
    elapsed_time = end_time - start_time

    # 4. 计算统计指标
    print(f"\n{'='*70}")
    print("📈 测试结果统计")
    print(f"{'='*70}")

    print(f"\n⏱️  总耗时: {elapsed_time:.2f}秒 ({elapsed_time/total_samples:.3f}秒/样本)")
    print(f"🚀 速度提升: {max_workers}倍并发")
    print(f"\n📊 API调用统计:")
    print(f"  ✅ 成功: {success_count} ({success_count/total_samples*100:.1f}%)")
    print(f"  ⛔ 规则层拦截: {rule_blocked_count} ({rule_blocked_count/total_samples*100:.1f}%)")
    print(f"  ❌ 失败: {error_count} ({error_count/total_samples*100:.1f}%)")

    # 只计算成功和规则拦截的样本准确率
    valid_samples = df[df['api_status'].isin(['success', 'rule_blocked'])]
    valid_count = len(valid_samples)

    if valid_count > 0:
        correct_count = valid_samples['correct'].sum()
        accuracy = correct_count / valid_count * 100

        print(f"\n🎯 验证准确率:")
        print(f"  有效样本: {valid_count}/{total_samples}")
        print(f"  正确预测: {correct_count}")
        print(f"  **总体准确率: {accuracy:.2f}%**")

        # 分类别统计
        safe_samples = valid_samples[valid_samples['label'] == 0]
        unsafe_samples = valid_samples[valid_samples['label'] == 1]

        if len(safe_samples) > 0:
            safe_acc = safe_samples['correct'].sum() / len(safe_samples) * 100
            print(f"  安全类(0)准确率: {safe_acc:.2f}% ({safe_samples['correct'].sum()}/{len(safe_samples)})")

        if len(unsafe_samples) > 0:
            unsafe_acc = unsafe_samples['correct'].sum() / len(unsafe_samples) * 100
            print(f"  危险类(1)准确率: {unsafe_acc:.2f}% ({unsafe_samples['correct'].sum()}/{len(unsafe_samples)})")

        # 混淆矩阵
        tp = len(valid_samples[(valid_samples['label'] == 1) & (valid_samples['predicted_label'] == 1)])
        tn = len(valid_samples[(valid_samples['label'] == 0) & (valid_samples['predicted_label'] == 0)])
        fp = len(valid_samples[(valid_samples['label'] == 0) & (valid_samples['predicted_label'] == 1)])
        fn = len(valid_samples[(valid_samples['label'] == 1) & (valid_samples['predicted_label'] == 0)])

        print(f"\n📊 混淆矩阵:")
        print(f"  真正例(TP): {tp}")
        print(f"  真负例(TN): {tn}")
        print(f"  假正例(FP): {fp}")
        print(f"  假负例(FN): {fn}")

        if tp + fp > 0:
            precision = tp / (tp + fp) * 100
            print(f"\n  精确率(Precision): {precision:.2f}%")

        if tp + fn > 0:
            recall = tp / (tp + fn) * 100
            print(f"  召回率(Recall): {recall:.2f}%")

        if precision + recall > 0:
            f1 = 2 * (precision * recall) / (precision + recall)
            print(f"  F1分数: {f1:.2f}%")

    # 5. 保存结果
    print(f"\n{'='*70}")
    print(f"💾 保存结果到: {output_csv}")

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_csv), exist_ok=True)

    df.to_csv(output_csv, index=False, encoding='utf-8-sig')

    # 保存统计摘要
    summary_file = output_csv.replace('.csv', '_summary.txt')
    with open(summary_file, 'w', encoding='utf-8') as f:
        f.write(f"批量API验证测试报告\n")
        f.write(f"{'='*70}\n")
        f.write(f"测试时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"数据集: {input_csv}\n")
        f.write(f"样本数: {total_samples}\n")
        f.write(f"阈值: {threshold}\n")
        f.write(f"\n")
        f.write(f"API调用统计:\n")
        f.write(f"  成功: {success_count} ({success_count/total_samples*100:.1f}%)\n")
        f.write(f"  规则层拦截: {rule_blocked_count} ({rule_blocked_count/total_samples*100:.1f}%)\n")
        f.write(f"  失败: {error_count} ({error_count/total_samples*100:.1f}%)\n")
        f.write(f"\n")
        if valid_count > 0:
            f.write(f"验证准确率:\n")
            f.write(f"  总体准确率: {accuracy:.2f}%\n")
            f.write(f"  正确预测: {correct_count}/{valid_count}\n")

    print(f"📄 统计摘要: {summary_file}")
    print(f"\n✅ 测试完成!")
    print(f"{'='*70}\n")


if __name__ == "__main__":
    import sys

    # 可以通过命令行参数指定最大样本数和并发数
    max_samples = None
    max_workers = 2  # 默认2并发,避免API压力过大

    if len(sys.argv) > 1:
        max_samples = int(sys.argv[1])
        print(f"⚠️  限制测试样本数: {max_samples}")

    if len(sys.argv) > 2:
        max_workers = int(sys.argv[2])
        print(f"⚡ 并发线程数: {max_workers}")

    # 获取脚本所在目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)  # safety_api目录

    batch_verify_dataset(
        input_csv=os.path.join(project_root, "data/verify.csv"),
        output_csv=os.path.join(script_dir, "results/verify_api_results.csv"),
        threshold=0.5,
        max_samples=max_samples,
        max_workers=max_workers
    )
